package com.lpssfxy.datainit

import com.lpssfxy.datainit.entities.{Product, Rating}
import com.lpssfxy.datainit.utils.AppUtils
import com.mongodb.client.MongoClients
import com.mongodb.client.model.{Filters, IndexOptions, Indexes, Updates}
import org.apache.spark.SparkConf
import org.apache.spark.sql.{DataFrame, SparkSession}

object DataInit {
  // 主程序的入口
  def main(args: Array[String]): Unit = {
    if (args.length < 2) {
      println("请传入产品数据文件路径和评分数据文件路径")
      System.exit(1)
    }
    val productPath = args(0)
    val ratingPath = args(1)

    // 创建一个SparkConf配置
    val sparkConf = new SparkConf().setAppName("DataInit")

    // 创建一个SparkSession
    val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    // 在对DataFrame和Dataset进行操作许多操作都需要这个包进行支持
    import spark.implicits._
    // 将Product、Rating数据集加载进来
    val productRDD = spark.sparkContext.textFile(productPath)
    //将ProdcutRDD装换为DataFrame
    val productDF = productRDD.map(item => {
      val attr = item.split("\\^")
      Product(attr(0).toInt, attr(1).trim, attr(4).trim, attr(5).trim, attr(6).trim)
    }).toDF()

    val ratingRDD = spark.sparkContext.textFile(ratingPath)
    //将ratingRDD转换为DataFrame
    val ratingDF = ratingRDD.map(item => {
      val attr = item.split(",")
      Rating(attr(0).toInt, attr(1).toInt, attr(2).toDouble, attr(3).toInt)
    }).toDF()

    // 将数据保存到MongoDB中
    storeDataInMongoDB(productDF, ratingDF)

    // 关闭Spark
    spark.stop()
  }

  private def storeDataInMongoDB(productDF: DataFrame, ratingDF: DataFrame): Unit = {
    // 将当前数据写入到 MongoDB
    productDF
      .write
      .option("uri", AppUtils.getMongoUri)
      .option("collection", AppUtils.MONGODB_PRODUCT_COLLECTION)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()

    // 创建索引
    val mongoClient = MongoClients.create(AppUtils.getMongoUri)
    val mongoDatabase = mongoClient.getDatabase(AppUtils.getMongoDb)
    val productCollection = mongoDatabase.getCollection(AppUtils.MONGODB_PRODUCT_COLLECTION)
    val ratingCollection = mongoDatabase.getCollection(AppUtils.MONGODB_RATING_COLLECTION)

    productCollection.createIndex(Indexes.ascending("productId"))
    // 创建唯一的复合索引
    val indexOptions = new IndexOptions().unique(true)
    ratingCollection.createIndex(
      Indexes.compoundIndex(Indexes.ascending("userId"), Indexes.ascending("productId")),
      indexOptions
    )

    // 处理评分数据，确保只记录最新的评分
    ratingDF.collect().foreach { row =>
      val userId = row.getAs[Int]("userId")
      val productId = row.getAs[Int]("productId")
      val score = row.getAs[Double]("score")
      val timestamp = row.getAs[Long]("timestamp")

      // 使用 updateOne 操作，结合 upsert 选项
      ratingCollection.updateOne(
        Filters.and(Filters.eq("userId", userId), Filters.eq("productId", productId)),
        Updates.combine(
          Updates.set("score", score),
          Updates.set("timestamp", timestamp)
        ),
        new com.mongodb.client.model.UpdateOptions().upsert(true)
      )
    }
    // 关闭 MongoDB 的连接
    mongoClient.close()
  }
}